import boto3

from framework.common.pt_logger import logger
from framework.core.crawler_task import Task
from framework.core.sqs_task_pool import SQSTaskPool
from jobs.amazon.util import sqs_util
from util.athenaUtil import QueryExecutor, convert_to_list

athena = boto3.client('athena')


def setup_task_pool(name):
    pool_name = "pt-spider-serverless-%s-queue" % name
    task_pool = SQSTaskPool(name=pool_name)
    sqs_util.create_if_not_existed(queue_name=pool_name)
    return task_pool


def start_initiator(metadata=None, payload=None):
    assert "tt" in payload and "athena_sql" in payload
    batch_id = metadata["time"].split("T")[0]
    task_type = payload["tt"]
    # JimmyMo sql to query asin and rank (asin must be at index 0)
    # e.g. athena_sql = f'SELECT "asin", "rank" FROM "coreproducts"."manual_testing_2021_07_01" where nodeid=\'12897221\''
    athena_sql = payload["athena_sql"]
    # node_id = payload["nodeId"]

    task_pool = setup_task_pool(task_type)
    count = sqs_util.check_sqs_remain_message_count(task_pool.queue_name)
    if count and int(count) <= 0:
        # query all asins grouped by one item_type, from athena and compose the init_tasks
        # athena_sql = f'SELECT * FROM "coreproducts"."manual_testing_2021_07_01" where nodeid=\'{node_id}\''
        athenaQuery = QueryExecutor()
        result = athenaQuery.adhoc_query(athena_sql)
        lst = convert_to_list(result)
        asin_list = [item[0] for idx, item in enumerate(lst) if idx != 0]
        url_template = "https://www.amazon.com/dp/{asin}"
        task_url_list = [url_template.format(asin=v) for v in asin_list]
        task_list = [(url, {"batchId": batch_id, "SpiderName": "SelectiveSpider", "tt": task_type}, None, None) for
                     url in task_url_list]
        init_tasks = Task.create_tasks(task_list, parser="dpParser", parent_result=None)
        task_pool.enqueue(init_tasks)
    else:
        logger.warning("there are still %s messages in the queue %s, please purge the queue and restart" % (count, task_pool.queue_name))


if __name__ == "__main__":
    batch_id = "2021-07-01"
    task_type = "lazyWidget"
    node_id = "12897221"
    # athena_sql = f'SELECT "asin", "rank" FROM "coreproducts"."manual_testing_2021_07_01" where nodeid=\'{node_id}\';"'
    athena_sql = f'SELECT "asin", "rank" FROM "coreproducts"."manual_testingttt_2021_07_01" where rawnodeid=\'{node_id}\''
    athenaQuery = QueryExecutor()
    result = athenaQuery.adhoc_query(athena_sql)
    lst = convert_to_list(result)
    asin_list = [item[2] for idx, item in enumerate(lst) if idx != 0]
    url_template = "https://www.amazon.com/hz/reviews-render/ajax/lazy-widgets/stream?asin={asin}&language=en_US&lazyWidget=cr-summarization-attributes&lazyWidget=cr-age-recommendation&lazyWidget=cr-solicitation&lazyWidget=cr-summarization-lighthut"
    task_url_list = [url_template.format(asin=v) for v in asin_list]
    task_list = [(url, {"batchId": batch_id, "SpiderName": "CoreProductsSpider", "tt": task_type}, None, None) for url in task_url_list]
    init_tasks = Task.create_tasks(task_list, parser="default", parent_result=None)
    print(result)
